import numpy as np

def softmax(x):
    if x.ndim == 2:
        x = x.T
        x = x - np.max(x,axis=0)
        y = np.exp(x)/np.sum(np.exp(x),axis=0)
        return y.T
    x = x - np.max(x)
    return np.exp(x)/np.sum(np.exp(x))

a = np.array([[0.3,2.9,4.0],[1,3,5]])
y = softmax(a)
print(y)